package tools

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
	"time"

	"workflow-gui/internal/ai"
	"workflow-gui/internal/database"
	"workflow-gui/internal/protocol"
	"workflow-gui/internal/workflow"
)

// AISupplementHandler handles AI supplement operations
type AISupplementHandler struct {
	db           *database.DB
	aiClient     *ai.AIClient
	promptEngine *ai.PromptEngine
}

// NewAISupplementHandler creates a new AI supplement handler
func NewAISupplementHandler(db *database.DB) (*AISupplementHandler, error) {
	// Create prompt engine
	promptEngine, err := ai.NewPromptEngine()
	if err != nil {
		return nil, fmt.Errorf("failed to create prompt engine: %w", err)
	}

	return &AISupplementHandler{
		db:           db,
		aiClient:     nil, // Not needed - Claude AI executes prompts
		promptEngine: promptEngine,
	}, nil
}

// HandleAISupplementStage1 handles Stage 1 AI supplement request
func (h *AISupplementHandler) HandleAISupplementStage1(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract change_id
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	// Load project from database
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	// Check workflow_stage (must be 'draft')
	if project.Meta.WorkflowStage != "draft" {
		return protocol.NewErrorResult(fmt.Sprintf("Invalid workflow_stage: %s (expected 'draft')", project.Meta.WorkflowStage)), nil
	}

	// Build Stage 1 prompt
	prompt, err := h.buildStage1Prompt(project)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to build prompt: %v", err)), nil
	}

	// Update workflow_stage to 'ai_reviewing_stage1'
	project.Meta.WorkflowStage = "ai_reviewing_stage1"
	project.Meta.UpdatedAt = time.Now()
	if err := dao.Update(changeID, project); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to update workflow_stage: %v", err)), nil
	}

	// Return prompt for Claude AI to execute
	message := fmt.Sprintf("📋 **Stage 1 补充提示词已生成**\n\n请根据以下提示词向用户提问，收集信息后补充 proposals 字段。\n\n---\n\n%s\n\n---\n\n**执行步骤**：\n1. 根据提示词向用户提问\n2. 收集用户回答\n3. 理解并补充字段（priority, target_users, key_features 等）\n4. 调用 `save_stage1_results` 工具保存结果",
		prompt)

	return protocol.NewSuccessResult(message), nil
}

// buildStage1Prompt builds the Stage 1 prompt
func (h *AISupplementHandler) buildStage1Prompt(project *workflow.WorkflowData) (string, error) {
	// Construct prompt data
	data := ai.PromptData{
		Summary:    project.Proposal.Summary,
		Motivation: project.Proposal.Motivation,
		ChangeType: project.Proposal.Scope.ChangeType,
	}

	// Render prompt
	prompt, err := h.promptEngine.RenderPrompt("stage1_supplement", data)
	if err != nil {
		return "", fmt.Errorf("failed to render prompt: %w", err)
	}

	return prompt, nil
}

// parseStage1Response parses AI response and extracts JSON
func (h *AISupplementHandler) parseStage1Response(response string) (map[string]interface{}, error) {
	// Extract JSON from response (AI might wrap it in markdown code blocks)
	jsonPattern := regexp.MustCompile("(?s)```json\\s*(.+?)\\s*```")
	matches := jsonPattern.FindStringSubmatch(response)

	var jsonStr string
	if len(matches) > 1 {
		jsonStr = matches[1]
	} else {
		// Try to find JSON object directly
		jsonStr = strings.TrimSpace(response)
		if !strings.HasPrefix(jsonStr, "{") {
			return nil, fmt.Errorf("no JSON found in AI response")
		}
	}

	// Parse JSON
	var results map[string]interface{}
	if err := json.Unmarshal([]byte(jsonStr), &results); err != nil {
		return nil, fmt.Errorf("failed to parse JSON: %w", err)
	}

	return results, nil
}

// HandleSaveStage1Results handles saving Stage 1 results
func (h *AISupplementHandler) HandleSaveStage1Results(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract change_id
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	// Extract results
	results, ok := args["results"].(map[string]interface{})
	if !ok {
		return protocol.NewErrorResult("Missing or invalid argument: results"), nil
	}

	// Validate output
	if err := ai.ValidateStage1Output(results); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Validation failed: %v", err)), nil
	}

	// Save results to database
	if err := h.saveStage1Results(changeID, results); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to save results: %v", err)), nil
	}

	// Update workflow_stage to 'stage1_confirmed'
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	project.Meta.WorkflowStage = "stage1_confirmed"
	project.Meta.UpdatedAt = time.Now()
	if err := dao.Update(changeID, project); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to update workflow_stage: %v", err)), nil
	}

	// Format success message
	message := fmt.Sprintf("✅ Stage 1 结果已保存！\n\n已补充字段：\n- priority: %v\n- target_users: %v\n- key_features: %v\n- affected_capabilities: %v\n- breaking_changes: %v\n\n状态已更新为: stage1_confirmed\n\n可以继续调用 `ai_supplement_stage2` 进入下一阶段。",
		results["priority"],
		results["target_users"],
		results["key_features"],
		results["affected_capabilities"],
		results["breaking_changes"])

	return protocol.NewSuccessResult(message), nil
}

// saveStage1Results saves Stage 1 results to database
func (h *AISupplementHandler) saveStage1Results(changeID string, results map[string]interface{}) error {
	// Start transaction
	tx, err := h.db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Get project ID
	var projectID int
	err = tx.QueryRow("SELECT id FROM projects WHERE change_id = ?", changeID).Scan(&projectID)
	if err != nil {
		return fmt.Errorf("failed to get project ID: %w", err)
	}

	// Extract fields from results
	priority, _ := results["priority"].(string)
	targetUsers, _ := results["target_users"].(string)
	breakingChanges, _ := results["breaking_changes"].(bool)

	// Build scope JSON (update existing scope with new fields)
	scopeData := map[string]interface{}{
		"change_type":           results["change_type"],
		"priority":              priority,
		"target_users":          targetUsers,
		"key_features":          results["key_features"],
		"affected_capabilities": results["affected_capabilities"],
		"breaking_changes":      breakingChanges,
		"affected_files":        results["affected_files_json"],
		"mcp_tools":             results["mcp_tools"],
		"output_format":         results["output_format"],
	}

	scopeJSON, err := json.Marshal(scopeData)
	if err != nil {
		return fmt.Errorf("failed to marshal scope: %w", err)
	}

	// Update proposals table
	_, err = tx.Exec(`
		UPDATE proposals
		SET priority = ?, breaking_changes = ?, scope_json = ?
		WHERE project_id = ?`,
		priority, breakingChanges, scopeJSON, projectID)
	if err != nil {
		return fmt.Errorf("failed to update proposals: %w", err)
	}

	// Save reasoning to project_history
	historyEntry := map[string]interface{}{
		"stage":                           "stage1",
		"action":                          "ai_supplement",
		"priority_reasoning":              results["priority_reasoning"],
		"target_users_reasoning":          results["target_users_reasoning"],
		"key_features_reasoning":          results["key_features_reasoning"],
		"affected_capabilities_reasoning": results["affected_capabilities_reasoning"],
		"breaking_changes_reasoning":      results["breaking_changes_reasoning"],
	}

	historyJSON, err := json.Marshal(historyEntry)
	if err != nil {
		return fmt.Errorf("failed to marshal history: %w", err)
	}

	_, err = tx.Exec(`
		INSERT INTO project_history (project_id, timestamp, event, details)
		VALUES (?, ?, ?, ?)`,
		projectID, time.Now(), "ai_supplement_stage1", string(historyJSON))
	if err != nil {
		return fmt.Errorf("failed to insert history: %w", err)
	}

	// Commit transaction
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}

// HandleAISupplementStage2 handles Stage 2 supplement request
func (h *AISupplementHandler) HandleAISupplementStage2(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract change_id
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	// Load project from database
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	// Check workflow_stage
	if project.Meta.WorkflowStage != "stage1_confirmed" {
		return protocol.NewErrorResult(fmt.Sprintf("Invalid workflow_stage: expected 'stage1_confirmed', got '%s'", project.Meta.WorkflowStage)), nil
	}

	// Build Stage 2 prompt
	prompt, err := h.buildStage2Prompt(project)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to build prompt: %v", err)), nil
	}

	// Update workflow_stage to 'ai_reviewing_stage2'
	project.Meta.WorkflowStage = "ai_reviewing_stage2"
	project.Meta.UpdatedAt = time.Now()
	if err := dao.Update(changeID, project); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to update workflow_stage: %v", err)), nil
	}

	// Return prompt for Claude AI to execute
	message := fmt.Sprintf("📋 **Stage 2 补充提示词已生成**\n\n请根据以下提示词生成 requirements（需求规范）。\n\n---\n\n%s\n\n---\n\n**执行步骤**：\n1. 根据 Stage 1 已确认的 proposals 生成 requirements\n2. 每个 requirement 必须包含：id, operation, name, description\n3. description 必须使用 SHALL 或 MUST 关键词\n4. 调用 `save_stage2_results` 工具保存结果",
		prompt)

	return protocol.NewSuccessResult(message), nil
}

// buildStage2Prompt builds the Stage 2 prompt
func (h *AISupplementHandler) buildStage2Prompt(project *workflow.WorkflowData) (string, error) {
	// Construct Stage 2 prompt data
	data := ai.Stage2PromptData{
		Summary:              project.Proposal.Summary,
		Motivation:           project.Proposal.Motivation,
		Priority:             project.Proposal.Scope.Priority,
		TargetUsers:          project.Proposal.Scope.TargetUsers,
		KeyFeatures:          project.Proposal.Scope.KeyFeatures,
		AffectedCapabilities: project.Proposal.Scope.AffectedCapabilities,
		BreakingChanges:      project.Proposal.Scope.BreakingChanges,
	}

	// Render prompt
	prompt, err := h.promptEngine.RenderPrompt("stage2_supplement", data)
	if err != nil {
		return "", fmt.Errorf("failed to render prompt: %w", err)
	}

	return prompt, nil
}

// parseStage2Response parses AI response and extracts requirements
func (h *AISupplementHandler) parseStage2Response(response string) ([]workflow.Requirement, error) {
	// Extract JSON from markdown code blocks
	re := regexp.MustCompile(`(?s)` + "```json\\s*(.+?)\\s*```")
	matches := re.FindStringSubmatch(response)

	var jsonStr string
	if len(matches) > 1 {
		jsonStr = matches[1]
	} else {
		jsonStr = response
	}

	// Parse JSON
	var result struct {
		Requirements []workflow.Requirement `json:"requirements"`
	}
	if err := json.Unmarshal([]byte(jsonStr), &result); err != nil {
		return nil, fmt.Errorf("failed to parse JSON: %w", err)
	}

	return result.Requirements, nil
}

// saveStage2Results saves Stage 2 results to database
func (h *AISupplementHandler) saveStage2Results(changeID string, requirements []workflow.Requirement) error {
	// Begin transaction
	tx, err := h.db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Get project ID
	var projectID int
	err = tx.QueryRow("SELECT id FROM projects WHERE change_id = ?", changeID).Scan(&projectID)
	if err != nil {
		return fmt.Errorf("failed to get project ID: %w", err)
	}

	// Insert requirements
	for _, req := range requirements {
		_, err := tx.Exec(`
			INSERT INTO requirements (project_id, req_id, operation, name, description)
			VALUES (?, ?, ?, ?, ?)`,
			projectID, req.ID, req.Operation, req.Name, req.Description)
		if err != nil {
			return fmt.Errorf("failed to insert requirement: %w", err)
		}
	}

	// Update workflow_stage
	_, err = tx.Exec(`
		UPDATE projects
		SET workflow_stage = 'ai_reviewing_stage2', updated_at = ?
		WHERE change_id = ?`,
		time.Now(), changeID)
	if err != nil {
		return fmt.Errorf("failed to update workflow_stage: %w", err)
	}

	// Commit transaction
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}

// HandleAISupplementStage3 handles Stage 3 supplement request
func (h *AISupplementHandler) HandleAISupplementStage3(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract change_id
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	// Load project from database
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	// Check workflow_stage
	if project.Meta.WorkflowStage != "stage2_confirmed" {
		return protocol.NewErrorResult(fmt.Sprintf("Invalid workflow_stage: expected 'stage2_confirmed', got '%s'", project.Meta.WorkflowStage)), nil
	}

	// Build Stage 3 prompt
	prompt, err := h.buildStage3Prompt(project.Spec.Requirements)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to build prompt: %v", err)), nil
	}

	// Update workflow_stage to 'ai_reviewing_stage3'
	project.Meta.WorkflowStage = "ai_reviewing_stage3"
	project.Meta.UpdatedAt = time.Now()
	if err := dao.Update(changeID, project); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to update workflow_stage: %v", err)), nil
	}

	// Return prompt for Claude AI to execute
	message := fmt.Sprintf("📋 **Stage 3 补充提示词已生成**\n\n请根据以下提示词生成 scenarios（验收场景）。\n\n---\n\n%s\n\n---\n\n**执行步骤**：\n1. 根据 Stage 2 已确认的 requirements 生成 scenarios\n2. 每个 scenario 必须包含：id, name, when, then\n3. 使用 WHEN...THEN 格式描述场景\n4. 调用 `save_stage3_results` 工具保存结果",
		prompt)

	return protocol.NewSuccessResult(message), nil
}

// buildStage3Prompt builds the Stage 3 prompt
func (h *AISupplementHandler) buildStage3Prompt(requirements []workflow.Requirement) (string, error) {
	// Convert workflow.Requirement to ai.Requirement
	aiRequirements := make([]ai.Requirement, len(requirements))
	for i, req := range requirements {
		aiRequirements[i] = ai.Requirement{
			ReqID:       req.ID,
			Name:        req.Name,
			Description: req.Description,
		}
	}

	// Construct Stage 3 prompt data
	data := ai.Stage3PromptData{
		Requirements: aiRequirements,
	}

	// Render prompt
	prompt, err := h.promptEngine.RenderPrompt("stage3_supplement", data)
	if err != nil {
		return "", fmt.Errorf("failed to render prompt: %w", err)
	}

	return prompt, nil
}

// parseStage3Response parses AI response and extracts scenarios
func (h *AISupplementHandler) parseStage3Response(response string) ([]workflow.Scenario, error) {
	// Extract JSON from markdown code blocks
	re := regexp.MustCompile(`(?s)` + "```json\\s*(.+?)\\s*```")
	matches := re.FindStringSubmatch(response)

	var jsonStr string
	if len(matches) > 1 {
		jsonStr = matches[1]
	} else {
		jsonStr = response
	}

	// Parse JSON
	var result struct {
		Scenarios []workflow.Scenario `json:"scenarios"`
	}
	if err := json.Unmarshal([]byte(jsonStr), &result); err != nil {
		return nil, fmt.Errorf("failed to parse JSON: %w", err)
	}

	return result.Scenarios, nil
}

// saveStage3Results saves Stage 3 results to database
func (h *AISupplementHandler) saveStage3Results(changeID string, scenarios []workflow.Scenario) error {
	// Begin transaction
	tx, err := h.db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Get project ID
	var projectID int
	err = tx.QueryRow("SELECT id FROM projects WHERE change_id = ?", changeID).Scan(&projectID)
	if err != nil {
		return fmt.Errorf("failed to get project ID: %w", err)
	}

	// Insert scenarios
	for _, scenario := range scenarios {
		// Get requirement ID from req_id
		var requirementID int
		err = tx.QueryRow("SELECT id FROM requirements WHERE project_id = ? LIMIT 1", projectID).Scan(&requirementID)
		if err != nil {
			return fmt.Errorf("failed to get requirement ID: %w", err)
		}

		_, err := tx.Exec(`
			INSERT INTO scenarios (requirement_id, scenario_id, name, when_condition, then_result)
			VALUES (?, ?, ?, ?, ?)`,
			requirementID, scenario.ID, scenario.Name, scenario.When, scenario.Then)
		if err != nil {
			return fmt.Errorf("failed to insert scenario: %w", err)
		}
	}

	// Update workflow_stage to completed
	_, err = tx.Exec(`
		UPDATE projects
		SET workflow_stage = 'completed', updated_at = ?
		WHERE change_id = ?`,
		time.Now(), changeID)
	if err != nil {
		return fmt.Errorf("failed to update workflow_stage: %w", err)
	}

	// Commit transaction
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}

// HandleSaveStage2Results handles saving Stage 2 results
func (h *AISupplementHandler) HandleSaveStage2Results(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract change_id
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	// Extract requirements
	requirementsRaw, ok := args["requirements"].([]interface{})
	if !ok {
		return protocol.NewErrorResult("Missing or invalid argument: requirements"), nil
	}

	// Convert to workflow.Requirement
	requirements := make([]workflow.Requirement, len(requirementsRaw))
	for i, reqRaw := range requirementsRaw {
		reqMap, ok := reqRaw.(map[string]interface{})
		if !ok {
			return protocol.NewErrorResult(fmt.Sprintf("Invalid requirement at index %d", i)), nil
		}

		requirements[i] = workflow.Requirement{
			ID:          reqMap["id"].(string),
			Operation:   reqMap["operation"].(string),
			Name:        reqMap["name"].(string),
			Description: reqMap["description"].(string),
		}
	}

	// Validate output
	if err := ai.ValidateStage2Output(requirements); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Validation failed: %v", err)), nil
	}

	// Save results
	if err := h.saveStage2Results(changeID, requirements); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to save results: %v", err)), nil
	}

	// Update workflow_stage to 'stage2_confirmed'
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	project.Meta.WorkflowStage = "stage2_confirmed"
	project.Meta.UpdatedAt = time.Now()
	if err := dao.Update(changeID, project); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to update workflow_stage: %v", err)), nil
	}

	return protocol.NewSuccessResult(fmt.Sprintf("✅ Stage 2 结果已保存！\n\n已生成 %d 个 requirements\n\n状态已更新为: stage2_confirmed\n\n可以继续调用 `ai_supplement_stage3` 进入下一阶段。", len(requirements))), nil
}

// HandleSaveStage3Results handles saving Stage 3 results
func (h *AISupplementHandler) HandleSaveStage3Results(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract change_id
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	// Extract scenarios
	scenariosRaw, ok := args["scenarios"].([]interface{})
	if !ok {
		return protocol.NewErrorResult("Missing or invalid argument: scenarios"), nil
	}

	// Convert to workflow.Scenario
	scenarios := make([]workflow.Scenario, len(scenariosRaw))
	for i, scRaw := range scenariosRaw {
		scMap, ok := scRaw.(map[string]interface{})
		if !ok {
			return protocol.NewErrorResult(fmt.Sprintf("Invalid scenario at index %d", i)), nil
		}

		scenarios[i] = workflow.Scenario{
			ID:   scMap["id"].(string),
			Name: scMap["name"].(string),
			When: scMap["when"].(string),
			Then: scMap["then"].(string),
		}
	}

	// Validate output
	if err := ai.ValidateStage3Output(scenarios); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Validation failed: %v", err)), nil
	}

	// Save results
	if err := h.saveStage3Results(changeID, scenarios); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to save results: %v", err)), nil
	}

	// Update workflow_stage to 'completed'
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	project.Meta.WorkflowStage = "completed"
	project.Meta.Status = "ready_to_execute"
	project.Meta.UpdatedAt = time.Now()
	if err := dao.Update(changeID, project); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to update workflow_stage: %v", err)), nil
	}

	return protocol.NewSuccessResult(fmt.Sprintf("✅ Stage 3 结果已保存！\n\n已生成 %d 个 scenarios\n\n状态已更新为: completed\n\n工作流已完成，可以开始执行任务。", len(scenarios))), nil
}
